from krita import (krita, InfoObject)
from math import sqrt, ceil, floor
import json
# manages the dialog that lets you set user preferences
# before applying the script

from pathlib import Path
# for path operations (who'd have guessed)

from PyQt5.QtWidgets import QWidget, QLabel, QMessageBox
# for debug messages


class SpritesheetExporter(object):

    def __init__(self):
        # user-defined variables
        self.exportName = "Spritesheet"
        self.defaultPath = Path.home().joinpath("spritesheetExportKritaTmp")
        # remember this is a Path, not a string, and as such
        # you can't do string operations on it (unless you convert it first):
        self.exportDir = Path.home()
        # this is a Path too. Trust me.
        self.spritesExportDir = self.defaultPath
        self.isDirectionHorizontal = True
        self.defaultTime = -1
        self.defaultSpace = 0
        self.rows = self.defaultSpace
        self.columns = self.defaultSpace
        self.start = self.defaultTime
        self.end = self.defaultTime
        self.forceNew = False
        self.removeTmp = True
        self.step = 1
        self.layersAsAnimation = False
        self.writeTextureAtlas = False
        self.layersList = []
        self.layersStates = []
        self.offLayers = 0

    def positionLayer(self, layer, layerIDNum, frameIDNum, width, height):
        if self.isDirectionHorizontal:
            layer.move(int(frameIDNum % self.columns) * int(width),
                       (int(layerIDNum)) * int(height))
        else:
            layer.move(int(layerIDNum) * int(width),
                       int(frameIDNum % self.rows) * int(height))

    def checkLayerEnd(self, layer, doc):
        if layer.visible():
            if layer.animated():
                frame = doc.fullClipRangeEndTime()
                while not (layer.hasKeyframeAtTime(frame)
                    or frame < 0):
                    frame -=1
                if self.end < frame:
                    self.end = frame
            if len(layer.childNodes()) != 0:
            # if it was a group layer
            # we also check its kids
                for kid in layer.childNodes():
                    self.checkLayerEnd(kid, doc)

    def checkLayerStart(self, layer, doc):
        if layer.visible():
            if layer.animated():
                frame = 0
                while not (layer.hasKeyframeAtTime(frame)
                    or frame > doc.fullClipRangeEndTime()):
                    frame +=1
                if self.start > frame:
                    self.start = frame
            if len(layer.childNodes()) != 0:
                # if it was a group layer
                # we also check its kids
                for kid in layer.childNodes():
                    self.checkLayerStart(kid, doc)

    # get actual animation duration
    def setStartEndFrames(self, layer, doc):
        # only from version 4.2.x on can we use hasKeyframeAtTime;
        # in earlier versions we just export from 0 to 100 as default
        ver = Application.version ()
        isNewVersion = (int(ver[0]) > 4 or ( int(ver[0]) == 4 and int(ver[2]) >= 2))
        self.start = self.defaultTime
        self.end = self.defaultTime
        if(isNewVersion):
            self.checkLayerEnd(layer, doc)
            self.start = self.end
            self.checkLayerStart(layer, doc)
        else:
            self.end = 100
            self.start = 0
                
    # - export all frames of the animation in a temporary folder as png
    # - create a new document of the right dimensions
    #   according to self.rows and self.columns
    # - position each exported frame in the new doc according to its name
    # - export the doc (aka the spritesheet)
    # - remove tmp folder if needed
    def export(self, debugging=False):

        def debugPrint(message, usingTerminal = True):
            if usingTerminal:
                print(message)
            else:
                QMessageBox.information(QWidget(), i18n("Debug info: "),
                i18n(message))

        def sheetExportPath(suffix=""):
            return self.exportDir.joinpath(self.exportName + suffix)

        def spritesExportPath(suffix=""):
            return self.spritesExportDir.joinpath(self.exportName + suffix)


        def spriteFrameExportPath(layNum, frameNum):
            return str(spritesExportPath(fileNum(layNum) + fileNum(frameNum) + ".png"))

        def fileNum(num):
            return "_" + str(num).zfill(3)

        def exportFrame(layNum, frameNum, doc, debugging = False):
            doc.waitForDone()
            imagePath = spriteFrameExportPath(layNum, frameNum)
            doc.exportImage(imagePath, InfoObject())
            if(debugging):
                debugPrint("exporting frame " + str(frameNum) + " at " + imagePath)
            
        def getLayerState(layer, debugging = False):
            if len(layer.childNodes()) != 0:
                # if it was a group layer
                # we also check its kids
                for kid in layer.childNodes():
                    getLayerState(kid, debugging)
            else:
                self.layersStates.append(layer.visible() and layer.animated())
                self.layersList.append(layer)
                if(not (layer.visible() and layer.animated())):
                    self.offLayers += 1
                if(debugging):
                    debugPrint("saving state " + str(layer.visible()) + 
                    " of layer " + str(layer))

        def exportAnim(layerIDNum, doc, debugging = False):
            layer = self.layersList[layerIDNum]
            self.setStartEndFrames(layer, doc)
            doc.setCurrentTime(self.start)
            if(debugging):
                ver = Application.version ()
                isNewVersion = (int(ver[0]) > 4 or ( int(ver[0]) == 4 and int(ver[2]) >= 2))
                if (isNewVersion):
                    debugPrint(
                    "animation Length: " +
                    str(doc.animationLength()) +
                    "; full clip start: " +
                    str(doc.fullClipRangeStartTime()) +
                    "; full clip end: " +
                    str(doc.fullClipRangeEndTime()))
                debugPrint("export start: " +
                str(self.start) +
                "; export end: " +
                str(self.end) +
                "; export length: " +
                str(self.end - self.start)
                )
            framesNum = ((self.end + 1) - self.start)/self.step
            frameIDNum = self.start
            # export frames
            while(doc.currentTime() <= self.end):
                exportFrame(layerIDNum, frameIDNum, doc, debugging)
                frameIDNum += self.step
                doc.setCurrentTime(frameIDNum) 
            return framesNum           

        if debugging:
            print("")
            debugPrint("Export spritesheet start.")

        # clearing lists in case the script is used several times 
        # without restarting krita
        self.layersList.clear()
        self.layersStates.clear()
        self.offLayers = 0

        addedFolder = False
        # create a temporary export directory for the individual sprites
        # if the user didn't set any
        if self.spritesExportDir == self.defaultPath:
            self.spritesExportDir = sheetExportPath("_sprites")

        if self.forceNew and self.spritesExportDir.exists():
            exportNum = 0

            parentPath = self.spritesExportDir.parent
            folder = str(self.spritesExportDir.parts[-1])

            def exportCandidate():
                return parentPath.joinpath(folder + str(exportNum))

            # in case the user has a folder with the exact same name
            # as my temporary one
            while exportCandidate().exists():
                exportNum += 1
            self.spritesExportDir = exportCandidate()

        # if forceNew, spritesExportDir's value is taken
        # from the user-set choices in the dialog

        # this will always be called if not forceNew
        # because it will always create a new export folder
        if not (self.spritesExportDir).exists():
            addedFolder = True
            (self.spritesExportDir).mkdir()

        # render animation in the sprites export folder
        doc = Krita.instance().activeDocument()
        doc.setBatchmode(True)  # so it won't show the export dialog window


        #########################################################
        layerIDNum = 0
        framesNum = 0
        layers = doc.topLevelNodes()
        for layer in layers:
            getLayerState(layer, debugging)
        layersNum = len(self.layersList)

        if debugging:
            debugPrint("total layer num " + str(layersNum))
        # hide all layers
        for layer in self.layersList:
            layer.setVisible(False)

        # export each visible layer
        while(layerIDNum < layersNum):
            if (self.layersStates[layerIDNum] and self.layersList[layerIDNum].animated()):
                self.layersList[layerIDNum].setVisible(True)
                # refresh the canvas
                doc.refreshProjection()
                tmpNum = exportAnim(layerIDNum, doc, debugging)
                if tmpNum > framesNum:
                    framesNum = tmpNum
                self.layersList[layerIDNum].setVisible(False)  
            layerIDNum += 1
        
        # restore layers state
        layerIDNum = 0
        while(layerIDNum < layersNum):
            self.layersList[layerIDNum].setVisible(self.layersStates[layerIDNum])
            if(debugging):
                debugPrint("restoring layer " + str(layerIDNum))
            layerIDNum += 1
        layerIDNum = 0

        # getting current document info
        # so we can copy it over to the new document
        width = doc.width()
        height = doc.height()
        col = doc.colorModel()
        depth = doc.colorDepth()
        profile = doc.colorProfile()
        res = doc.resolution()

        self.columns = ceil(framesNum)
        self.rows = ceil(float(layersNum - self.offLayers))

        # creating a new document where we'll put our sprites
        sheet = Krita.instance().createDocument(
            self.columns * width,
            self.rows * height,
            self.exportName,
            col, depth, profile, res)
        if (debugging):
            debugPrint("new doc name: " + sheet.name())
            debugPrint("old doc width: " + str(width))
            debugPrint("num of columns: " + str(self.columns))
            debugPrint("num of rows: " + str(self.rows))
            debugPrint("new doc width: " + str(sheet.width()))

        # adding our sprites to the new document
        # and moving them to the right position
        root_node = sheet.rootNode()
        root_node.childNodes()[0].setVisible(False) # hide the default layer filled with white

        textureAtlas = { "frames": [ ] }

        layerNum = 0

        while (layerIDNum < layersNum):
            
            doc.waitForDone()
            
            if(self.layersStates[layerIDNum] and self.layersList[layerIDNum].animated()):

                self.setStartEndFrames(self.layersList[layerIDNum], doc)
                
                if(debugging):
                    debugPrint("managing start: " + str(self.start) + " end: " + str(self.end) + " layer: " + str(layerIDNum))
                
                frameIDNum = self.start

                while(frameIDNum <= self.end):
                    
                    doc.waitForDone()
                    
                    img = spriteFrameExportPath(layerIDNum, frameIDNum)
                    if(debugging):
                        debugPrint("managing file " + str(layerIDNum) + fileNum(frameIDNum) + " at " + img)
                    layer = sheet.createFileLayer(img, img, "ImageToSize")
                    root_node.addChildNode(layer, None)
                    self.positionLayer(
                        layer,
                        layerNum,
                        frameIDNum,
                        width,
                        height)
                    # refresh canvas so the layers actually do show
                    sheet.refreshProjection()
                    
                    if (self.removeTmp):
                        Path(img).unlink()
                    
                    if (self.writeTextureAtlas):
                        textureAtlas["frames"].append({
                            "filename": str(layerIDNum) + fileNum(frameIDNum),
                            "frame": { "x": layer.position().x(),
                                    "y": layer.position().y(),
                                    "w": width,
                                    "h": height}})
                    frameIDNum += 1
                
                if frameIDNum > self.end:
                    layerNum += 1
            
            layerIDNum += 1

        # export the document to the export location
        sheet.setBatchmode(True)  # so it won't show the export dialog window
        if debugging:
            debugPrint("exporting spritesheet to " + str(sheetExportPath()))

        sheet.exportImage(str(sheetExportPath(".png")), InfoObject())

        if (self.writeTextureAtlas):
            with open(str(sheetExportPath(".json")), 'w') as f:
                json.dump(textureAtlas, f)

        # and remove the empty tmp folder when you're done
        if self.removeTmp:
            if addedFolder:
                if debugging:
                    debugPrint("delete tmp dir")
                self.spritesExportDir.rmdir()

        if debugging:
            debugPrint("All done!")